#!/usr/bin/env python3
"""
使用示例：展示 arch_parser_optimized 的所有功能

这个示例展示了如何使用恢复后的所有功能，包括详细分析、过滤、分组等。
"""

import torch
import torch.nn as nn
from examples.arch_parser_optimized import (
    # 主要分析函数
    analyze_model,
    analyze_model_detailed,
    quick_model_summary,

    # 专用分析
    analyze_continual_learning_model,
    print_compact_detailed_analysis,

    # 导出功能
    export_model_analysis,

    # 遗留兼容性
    count_parameters,
    print_model_parameter_summary,

    # 直接使用分析器
    ModelAnalyzer,
    AnalysisType
)
import logging
# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format='%(name)s - %(levelname)s - %(message)s'
)

def create_example_model():
    """创建一个示例模型用于演示"""
    class ExampleContinualLearningModel(nn.Module):
        def __init__(self):
            super().__init__()
            # 主干网络 (通常在持续学习中冻结)
            self.backbone = nn.Sequential(
                nn.Conv2d(3, 32, 3, padding=1),
                nn.BatchNorm2d(32),
                nn.ReLU(),
                nn.Conv2d(32, 64, 3, padding=1),
                nn.BatchNorm2d(64),
                nn.AdaptiveAvgPool2d((4, 4))
            )

            # 分类头
            self.classifier = nn.Linear(64 * 4 * 4, 10)

            # 适配器层 (通常可训练)
            self.adapter = nn.Sequential(
                nn.Linear(64 * 4 * 4, 128),
                nn.ReLU(),
                nn.Linear(128, 64 * 4 * 4)
            )

            # 冻结主干网络
            for param in self.backbone.parameters():
                param.requires_grad = False

        def forward(self, x):
            features = self.backbone(x).flatten(1)
            adapted_features = features + self.adapter(features)
            return self.classifier(adapted_features)

    return ExampleContinualLearningModel()


def demo_basic_analysis():
    """演示基础分析功能"""
    print("="*80)
    print("基础分析功能演示")
    print("="*80)

    model = create_example_model()

    # 1. 层次化分析 (默认)
    print("\n1. 层次化分析:")
    analyze_model(model, "示例模型", "hierarchical")

    # 2. 超紧凑摘要
    print("\n2. 超紧凑摘要:")
    analyze_model(model, "示例模型", "ultra_compact")

    # 3. 紧凑详细分析
    print("\n3. 紧凑详细分析:")
    analyze_model(model, "示例模型", "compact_detailed")

    # 4. 按组件分组分析
    print("\n4. 按组件分组分析:")
    analyze_model(model, "示例模型", "component_grouped")


def demo_detailed_analysis():
    """演示详细分析功能"""
    print("="*80)
    print("详细分析功能演示")
    print("="*80)

    model = create_example_model()

    # 1. 默认详细分析
    print("\n1. 默认详细分析 (层次化):")
    analyze_model_detailed(model, "详细分析示例")

    # 2. 仅显示可训练参数
    print("\n2. 仅显示可训练参数:")
    analyze_model_detailed(
        model,
        "可训练参数分析",
        trainable_only=True
    )

    # 3. 显示前10个参数
    print("\n3. 显示前10个参数:")
    analyze_model_detailed(
        model,
        "Top 10 参数",
        show_top_n=10,
        sort_by="total_params"
    )

    # 4. 按组件分组 (非层次化)
    print("\n4. 按组件分组显示:")
    analyze_model_detailed(
        model,
        "按组件分组",
        group_by="component",
        hierarchical=False
    )

    # 5. 按层类型分组
    print("\n5. 按层类型分组:")
    analyze_model_detailed(
        model,
        "按层类型分组",
        group_by="layer_type",
        hierarchical=False
    )


def demo_filtering_options():
    """演示参数过滤选项"""
    print("="*80)
    print("参数过滤选项演示")
    print("="*80)

    model = create_example_model()
    analyzer = ModelAnalyzer(model, "过滤示例")

    # 获取所有参数
    all_params = analyzer.param_details
    print(f"总参数数量: {len(all_params)}")

    # 过滤选项演示
    trainable_params = analyzer.filter_parameters(trainable_only=True)
    print(f"可训练参数数量: {len(trainable_params)}")

    shallow_params = analyzer.filter_parameters(max_depth=3)
    print(f"最大深度1的参数: {len(shallow_params)}")

    backbone_params = analyzer.filter_parameters(component_filter="backbone")
    print(f"主干网络参数: {len(backbone_params)}")

    # 组合过滤
    filtered = analyzer.filter_parameters(
        trainable_only=True,
        component_filter="classifier"
    )
    print(f"可训练的分类器参数: {len(filtered)}")


def demo_component_statistics():
    """演示组件统计功能"""
    print("="*80)
    print("组件统计功能演示")
    print("="*80)

    model = create_example_model()
    analyzer = ModelAnalyzer(model, "组件统计示例")

    # 获取组件统计
    component_stats = analyzer.get_parameter_statistics_by_component()

    print("各组件详细统计:")
    for component, stats in component_stats.items():
        print(f"\n{component.upper()} 组件:")
        print(f"  总参数: {stats['total_parameters']:,}")
        print(f"  可训练: {stats['trainable_parameters']:,}")
        print(f"  冻结: {stats['frozen_parameters']:,}")
        print(f"  层数: {stats['layer_count']}")
        print(f"  内存: {stats['memory_mb']:.2f} MB")
        print(f"  可训练比例: {stats['trainable_ratio']:.1%}")
        print(f"  参数个数: {stats['parameter_count']}")


def demo_export_functionality():
    """演示导出功能"""
    print("="*80)
    print("导出功能演示")
    print("="*80)

    model = create_example_model()

    # 1. 通过公共API导出
    print("1. 通过API导出:")
    export_model_analysis(model, "example_export_api.csv", "API导出示例")

    # 2. 通过分析器直接导出
    print("2. 通过分析器导出:")
    analyzer = ModelAnalyzer(model, "直接导出示例")
    analyzer.export_to_csv("example_export_direct.csv")

    # 验证导出文件
    import os
    for filename in ["example_export_api.csv", "example_export_direct.csv"]:
        if os.path.exists(filename):
            size = os.path.getsize(filename)
            print(f"✓ {filename}: {size} 字节")


def demo_legacy_compatibility():
    """演示遗留兼容性功能"""
    print("="*80)
    print("遗留兼容性功能演示")
    print("="*80)

    model = create_example_model()

    # 1. 参数计数
    total = count_parameters(model)
    trainable = count_parameters(model, trainable=True)
    print(f"参数计数: 总数={total:,}, 可训练={trainable:,}")

    # 2. 参数摘要 (遗留格式)
    print("\n参数摘要 (遗留格式):")
    summary = print_model_parameter_summary(model, task_id=1)

    # 3. 快速摘要
    print("\n快速摘要:")
    quick_summary = quick_model_summary(model, task_id=1)

    # 4. 持续学习模型专用分析
    print("\n持续学习模型分析:")
    cl_analysis = analyze_continual_learning_model(model, task_id=1)

    # 5. 紧凑详细分析 (函数版本)
    print("\n紧凑详细分析:")
    print_compact_detailed_analysis(model, "遗留兼容示例", max_components=5)


def demo_advanced_usage():
    """演示高级用法"""
    print("="*80)
    print("高级用法演示")
    print("="*80)

    model = create_example_model()

    # 1. 直接使用分析器类
    analyzer = ModelAnalyzer(model, "高级用法示例")

    # 2. 使用不同的显示策略
    print("\n使用策略模式:")

    # 详细分析策略，带自定义参数
    analyzer.display_analysis(
        AnalysisType.DETAILED,
        sort_by="total_params",
        show_top_n=5,
        trainable_only=True
    )

    # 3. 自定义详细参数表格
    print("\n自定义参数表格:")
    analyzer.print_detailed_parameter_table(
        sort_by="memory_mb",
        show_top_n=8,
        hierarchical=True
    )

    # 4. 获取原始数据进行自定义处理
    print("\n获取原始分析数据:")
    basic_stats = analyzer.basic_stats
    param_details = analyzer.param_details
    component_stats = analyzer.component_stats
    memory_analysis = analyzer.memory_analysis

    print(f"基础统计: {basic_stats}")
    print(f"参数详情数量: {len(param_details)}")
    print(f"组件数量: {len(component_stats)}")
    print(f"内存分析: {memory_analysis}")


def main():
    """运行所有演示"""
    print("ARCH_PARSER_OPTIMIZED 功能演示")
    print("="*80)
    print("这个示例展示了所有恢复和增强的功能")

    # 运行所有演示
    demo_basic_analysis()
    demo_detailed_analysis()
    demo_filtering_options()
    demo_component_statistics()
    demo_export_functionality()
    demo_legacy_compatibility()
    demo_advanced_usage()

    print("="*80)
    print("演示完成!")
    print("="*80)
    print("所有功能都已成功恢复并增强。")
    print("现在 arch_parser_optimized.py 包含了原版本的所有功能，")
    print("同时具有更好的代码设计和扩展性。")


if __name__ == "__main__":
    main()
